import pandas as pd
import numpy as np

def extract_time_features(df, timestamp_col='timestamp'):
    """
    从时间戳中提取各种时间特征
    
    参数:
    df: pandas DataFrame，包含时间戳列
    timestamp_col: 时间戳列的列名，默认为'timestamp'
    
    返回:
    df: 添加了时间特征的DataFrame
    """
    # 确保时间戳列是datetime类型
    if not pd.api.types.is_datetime64_any_dtype(df[timestamp_col]):
        df[timestamp_col] = pd.to_datetime(df[timestamp_col])
    
    # 提取基本时间特征
    df['year'] = df[timestamp_col].dt.year
    df['month'] = df[timestamp_col].dt.month
    df['day'] = df[timestamp_col].dt.day
    df['hour'] = df[timestamp_col].dt.hour
    df['minute'] = df[timestamp_col].dt.minute
    df['second'] = df[timestamp_col].dt.second
    
    # 提取星期相关特征
    df['dayofweek'] = df[timestamp_col].dt.dayofweek  # 0=星期一, 6=星期日
    df['weekday'] = df[timestamp_col].dt.weekday  # 与dayofweek相同
    df['is_weekend'] = df['dayofweek'].apply(lambda x: 1 if x >= 5 else 0)  # 5=星期六, 6=星期日
    
    # 提取一年中的第几天、第几周
    df['dayofyear'] = df[timestamp_col].dt.dayofyear
    df['weekofyear'] = df[timestamp_col].dt.isocalendar().week.astype(int)  # 一年中的第几周
    
    # 提取月份相关特征
    df['is_month_start'] = df[timestamp_col].dt.is_month_start.astype(int)  # 是否为月初
    df['is_month_end'] = df[timestamp_col].dt.is_month_end.astype(int)      # 是否为月末
    
    # 提取季度相关特征
    df['quarter'] = df[timestamp_col].dt.quarter  # 季度
    df['is_quarter_start'] = df[timestamp_col].dt.is_quarter_start.astype(int)
    df['is_quarter_end'] = df[timestamp_col].dt.is_quarter_end.astype(int)
    
    # 提取年份相关特征
    df['is_year_start'] = df[timestamp_col].dt.is_year_start.astype(int)  # 是否为年初
    df['is_year_end'] = df[timestamp_col].dt.is_year_end.astype(int)      # 是否为年末
    
    # 提取时间周期特征（用于捕获周期性）
    df['hour_sin'] = np.sin(2 * np.pi * df['hour'] / 24)
    df['hour_cos'] = np.cos(2 * np.pi * df['hour'] / 24)
    df['dayofweek_sin'] = np.sin(2 * np.pi * df['dayofweek'] / 7)
    df['dayofweek_cos'] = np.cos(2 * np.pi * df['dayofweek'] / 7)
    df['month_sin'] = np.sin(2 * np.pi * df['month'] / 12)
    df['month_cos'] = np.cos(2 * np.pi * df['month'] / 12)
    
    return df

# 示例用法
if __name__ == "__main__":
    # 创建示例数据
    dates = pd.date_range(start='2023-01-01', end='2023-12-31', freq='M')
    df = pd.DataFrame({'timestamp': dates, 'value': np.random.randn(len(dates))})
    
    # 提取时间特征
    df_with_features = extract_time_features(df)
    
    # 打印前几行查看结果
    print("原始数据:")
    print(df.head())
    print("\n添加时间特征后的数据:")
    print(df_with_features.head())
    
    # 打印所有特征列名
    print("\n提取的时间特征列:")
    print([col for col in df_with_features.columns if col not in ['timestamp', 'value']])
